package git.soulbgm.health;

import git.soulbgm.pojo.Backend;
import git.soulbgm.config.ProxyConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.SocketException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 后端服务器健康检查器
 *
 * <p>定期检查后端服务器的健康状态，确保负载均衡器只将流量转发到健康的服务器。
 * 通过发送心跳包或探测请求来判断服务器是否可用。</p>
 *
 * <p>主要功能：</p>
 * <ul>
 *   <li>定期执行健康检查</li>
 *   <li>维护服务器健康状态</li>
 *   <li>提供健康状态查询接口</li>
 *   <li>支持自定义健康检查策略</li>
 * </ul>
 *
 * @author soulbgm
 * @version 1.0
 * @since 1.0
 */
public class BackendHealthChecker {
    private static final Logger logger = LoggerFactory.getLogger(BackendHealthChecker.class);

    private final Map<Backend, Integer> missedHeartbeats = new ConcurrentHashMap<>();
    private final Map<Backend, Boolean> healthStatus = new ConcurrentHashMap<>();
    private final ScheduledExecutorService scheduler;
    private final DatagramSocket heartbeatSocket;
    private static final byte[] HEARTBEAT_DATA = "HEARTBEAT".getBytes();
    private static final int MAX_MISSED_HEARTBEATS = 2; // 最大允许丢失心跳次数
    private static final int HEARTBEAT_TIMEOUT = 20; // 心跳超时时间(ms)

    public BackendHealthChecker(ScheduledExecutorService scheduler) throws SocketException {
        this.scheduler = scheduler;
        this.heartbeatSocket = new DatagramSocket();
        this.heartbeatSocket.setSoTimeout(HEARTBEAT_TIMEOUT);

        // 初始化所有后端服务器状态
        for (Backend backend : ProxyConfig.getInstance().getBackends()) {
            missedHeartbeats.put(backend, 0);
            healthStatus.put(backend, true);
        }
    }

    /**
     * 启动心跳检查
     */
    public void start() {
        // 每30ms执行一次心跳检查
        scheduler.scheduleWithFixedDelay(this::checkHeartbeats, 0, 30, TimeUnit.MILLISECONDS);
    }

    /**
     * 检查所有后端服务器的心跳
     */
    private void checkHeartbeats() {
        for (Backend backend : ProxyConfig.getInstance().getBackends()) {
            try {
                sendHeartbeat(backend);
                // 如果成功收到响应
                missedHeartbeats.put(backend, 0);
                if (!healthStatus.getOrDefault(backend, true)) {
                    healthStatus.put(backend, true);
                    logger.info("后端服务器恢复在线: {}:{}", backend.getHost(), backend.getPort());
                }
            } catch (IOException e) {
                // 心跳失败，增加计数
                int missed = missedHeartbeats.compute(backend, (k, v) -> (v == null ? 0 : v) + 1);
                if (missed >= MAX_MISSED_HEARTBEATS && healthStatus.getOrDefault(backend, true)) {
                    healthStatus.put(backend, false);
                    logger.error("后端服务器离线: {}:{}", backend.getHost(), backend.getPort());
                }
            }
        }
    }

    /**
     * 发送心跳并等待响应
     */
    private void sendHeartbeat(Backend backend) throws IOException {
        // 发送心跳包
        DatagramPacket heartbeatPacket = new DatagramPacket(
                HEARTBEAT_DATA,
                HEARTBEAT_DATA.length,
                InetAddress.getByName(backend.getHost()),
                backend.getPort()
        );
        heartbeatSocket.send(heartbeatPacket);

        // 等待响应
        byte[] receiveBuffer = new byte[HEARTBEAT_DATA.length];
        DatagramPacket receivePacket = new DatagramPacket(receiveBuffer, receiveBuffer.length);
        heartbeatSocket.receive(receivePacket);
    }

    /**
     * 获取后端服务器健康状态
     */
    public boolean isHealthy(Backend backend) {
        return healthStatus.getOrDefault(backend, false);
    }

    /**
     * 更新配置时重置健康状态
     */
    public void resetHealth() {
        missedHeartbeats.clear();
        healthStatus.clear();
        for (Backend backend : ProxyConfig.getInstance().getBackends()) {
            missedHeartbeats.put(backend, 0);
            healthStatus.put(backend, true);
        }
    }

    /**
     * 关闭心跳检查
     */
    public void stop() {
        if (heartbeatSocket != null && !heartbeatSocket.isClosed()) {
            heartbeatSocket.close();
        }
    }
} 